package com.wuwei.elearning.config.factory;

import com.wuwei.elearning.config.AiConfig;
import com.wuwei.elearning.question.model.ChatMessage;
import com.wuwei.elearning.question.model.ChatResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.web.client.HttpStatusCodeException;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * OpenAI客户端实现
 */
@Slf4j
public class OpenAiClient implements AiClient {

    private final AiConfig aiConfig;
    private final RestTemplate restTemplate;
    private static final Pattern DATA_PATTERN = Pattern.compile("data: (.*)");
    
    public OpenAiClient(AiConfig aiConfig) {
        this.aiConfig = aiConfig;
        this.restTemplate = new RestTemplate();
    }
    
    @Override
    public String streamChatCompletion(List<ChatMessage> messages, SseEmitter emitter, StringBuilder responseBuilder) {
        try {
            // 将消息转换为OpenAI格式
            List<Map<String, String>> openaiMessages = convertToOpenAiFormat(messages);
            
            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(openaiMessages, true);
            
            // 设置请求头
            HttpHeaders headers = createHeaders();
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
            
            log.info("发送流式请求到API");
            
            // 获取基本URL
            String url = aiConfig.getEndpoint();
            
            log.info("Stream API endpoint URL: {}", url);
            log.info("Request headers: {}", headers);
            log.info("Request body: {}", requestBody);
            
            try {
                // 使用ResponseExtractor处理流式响应
                restTemplate.execute(url, org.springframework.http.HttpMethod.POST, 
                    req -> {
                        req.getHeaders().addAll(headers);
                        req.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                        try {
                            // 使用ObjectMapper并明确设置UTF-8编码
                            com.fasterxml.jackson.databind.ObjectMapper mapper = new com.fasterxml.jackson.databind.ObjectMapper();
                            byte[] jsonBytes = mapper.writeValueAsBytes(requestBody);
                            req.getBody().write(jsonBytes);
                        } catch (IOException e) {
                            log.error("写入请求体时出错", e);
                        }
                    },
                    new ResponseExtractor<Void>() {
                        @Override
                        public Void extractData(ClientHttpResponse response) throws IOException {
                            try (BufferedReader reader = new BufferedReader(
                                    new InputStreamReader(response.getBody(), "UTF-8"))) {
                                String line;
                                while ((line = reader.readLine()) != null) {
                                    if (line.isEmpty()) continue;
                                    
                                    if (line.startsWith("data:")) {
                                        Matcher matcher = DATA_PATTERN.matcher(line);
                                        if (matcher.find()) {
                                            String data = matcher.group(1).trim();
                                            if ("[DONE]".equals(data)) {
                                                break;
                                            }
                                            
                                            try {
                                                // 解析JSON响应
                                                com.fasterxml.jackson.databind.ObjectMapper mapper = 
                                                    new com.fasterxml.jackson.databind.ObjectMapper();
                                                Map<String, Object> responseMap = mapper.readValue(data, Map.class);
                                                Map<String, Object> choices = (Map<String, Object>) ((List<?>) responseMap.get("choices")).get(0);
                                                Map<String, Object> delta = (Map<String, Object>) choices.get("delta");
                                                
                                                if (delta.containsKey("content")) {
                                                    String content = (String) delta.get("content");
                                                    responseBuilder.append(content);
                                                    
                                                    ChatResponse partialResponse = new ChatResponse();
                                                    partialResponse.setContent(responseBuilder.toString());
                                                    partialResponse.setIsComplete(false);
                                                    emitter.send(partialResponse);
                                                }
                                            } catch (Exception e) {
                                                log.error("解析流式响应时出错: {}", e.getMessage(), e);
                                            }
                                        }
                                    }
                                }
                            } catch (Exception e) {
                                log.error("处理流式响应时出错", e);
                                ChatResponse errorResponse = new ChatResponse();
                                errorResponse.setContent("处理API响应时出错: " + e.getMessage());
                                errorResponse.setIsComplete(true);
                                try {
                                    emitter.send(errorResponse);
                                } catch (IOException ex) {
                                    log.error("发送错误响应时出错", ex);
                                }
                            }
                            return null;
                        }
                    });
                
                log.info("API流式请求处理完成");
                
            } catch (Exception e) {
                log.error("API请求失败: {}", e.getMessage(), e);
                throw e;
            }
            
            return null;
            
        } catch (Exception e) {
            log.error("调用API时出错", e);
            return "错误：" + e.getMessage();
        }
    }
    
    
    
    @Override
    public String getModelId() {
        return aiConfig.getModelId();
    }
    
    /**
     * 将内部消息格式转换为OpenAI API所需格式
     */
    private List<Map<String, String>> convertToOpenAiFormat(List<ChatMessage> messages) {
        return messages.stream()
                .map(msg -> {
                    Map<String, String> map = new HashMap<>();
                    map.put("role", msg.getRole());
                    map.put("content", msg.getContent());
                    return map;
                })
                .collect(Collectors.toList());
    }
    
    /**
     * 创建请求体
     */
    private Map<String, Object> createRequestBody(List<Map<String, String>> messages, boolean stream) {
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", aiConfig.getModelId());
        
        // 确保消息内容是有效的UTF-8编码
        List<Map<String, String>> sanitizedMessages = messages.stream()
            .map(msg -> {
                Map<String, String> sanitizedMsg = new HashMap<>();
                sanitizedMsg.put("role", msg.get("role"));
                
                // 处理内容以确保UTF-8编码有效
                String content = msg.get("content");
                if (content != null) {
                    try {
                        // 重新编码内容以确保有效的UTF-8
                        byte[] bytes = content.getBytes("UTF-8");
                        content = new String(bytes, "UTF-8");
                    } catch (Exception e) {
                        log.warn("消息内容UTF-8编码处理失败，尝试移除无效字符", e);
                        // 移除可能导致问题的非法UTF-8字符
                        content = content.replaceAll("[^\\p{ASCII}]", "");
                    }
                }
                sanitizedMsg.put("content", content);
                return sanitizedMsg;
            })
            .collect(Collectors.toList());
            
        requestBody.put("messages", sanitizedMessages);
        requestBody.put("max_tokens", aiConfig.getMaxTokens());
        requestBody.put("temperature", aiConfig.getTemperature());
        requestBody.put("stream", stream);
        return requestBody;
    }

    
    /**
     * 创建HTTP请求头
     */
    private HttpHeaders createHeaders() {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        
        String endpoint = aiConfig.getEndpoint().toLowerCase();
        
        // 为不同的API端点使用不同的授权头设置
        if (endpoint.contains("ark.cn-beijing.volces.com")) {
            // 火山引擎API
            headers.set("Authorization", "Bearer " + aiConfig.getApiKey());
            log.info("使用火山引擎授权头");
        } else if (endpoint.contains("dashscope.aliyuncs.com")) {
            // 阿里云DashScope
            headers.set("Authorization", "Bearer " + aiConfig.getApiKey());
            log.info("使用阿里云Dashscope授权头");
        } else {
            // 标准OpenAI格式
            headers.setBearerAuth(aiConfig.getApiKey());
        }
        
        if (aiConfig.getOpenaiOrganization() != null && !aiConfig.getOpenaiOrganization().isEmpty()) {
            headers.set("OpenAI-Organization", aiConfig.getOpenaiOrganization());
        }
        
        return headers;
    }
} 